import numpy as np
class NaiveBayes:
    def fit(self, X, y):
        self.classes = np.unique(y)
        self.class_prior = {}
        self.conditional_prob = {}

        for c in self.classes:
            X_c = X[y == c]
            self.class_prior[c] = len(X_c) / len(X)
            self.conditional_prob[c] = {
                "mean": X_c.mean(axis=0),
                "std": X_c.std(axis=0) + 1e-8
            }

    def _gaussian_prob(self, x, mean, std):
        exponent = np.exp(-((x - mean) ** 2) / (2 * (std ** 2)))
        return (1 / (np.sqrt(2 * np.pi) * std)) * exponent

    def predict(self, X):
        preds = []
        for x in X:
            class_probs = {}
            for c in self.classes:
                prior = np.log(self.class_prior[c])
                cond_prob = np.sum(
                    np.log(self._gaussian_prob(x, self.conditional_prob[c]["mean"], self.conditional_prob[c]["std"])))
                class_probs[c] = prior + cond_prob
            preds.append(max(class_probs, key=class_probs.get))
        return preds


# 生成测试数据
X = np.random.randn(100, 2)
y = np.random.randint(0, 2, 100)

# 训练模型
nb = NaiveBayes()
nb.fit(X, y)

# 测试模型
X_test = np.random.randn(10, 2)
y_pred = nb.predict(X_test)

# 打印预测结果
print("预测结果：", y_pred)
